package com.tlzh.controller;


import com.fasterxml.jackson.databind.ObjectMapper;
import com.tlzh.pojo.dto.ChatRequest;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@RequestMapping("/chatV3")
@CrossOrigin
public class ChatV3Controller {

    @Resource
    private  WebClient webClient;


    private final ObjectMapper objectMapper = new ObjectMapper();
    private static final Logger logger = LoggerFactory.getLogger(FluxController.class);

    private final Map<String, Disposable> activeRequests = new ConcurrentHashMap<>();

    // 存储每个用户的中断状态
    private  Map<String, AtomicBoolean> interruptionStatusMapV2 = new ConcurrentHashMap<>();


    @PostMapping(value = "/ai/ollamaStreamV1", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> queryModelV1(@RequestBody ChatRequest request) {
        String uuid = request.getToken();
        try {
            String requestBody = objectMapper.writeValueAsString(request);
            logger.info("Sending request body: {}", requestBody);
        } catch (Exception e) {
            logger.error("Error converting request to JSON", e);
        }

        // 初始化用户的中断状态
        AtomicBoolean isInterrupted = interruptionStatusMapV2.computeIfAbsent(uuid, k -> new AtomicBoolean(false));

        Flux<ServerSentEvent<String>> flux = webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .header("Accept-Charset", "UTF-8")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class)
                .map(response -> ServerSentEvent.<String>builder().data(response).build())
                .takeWhile(part -> !isInterrupted.get())
                .doFinally(signal -> {
                    interruptionStatusMapV2.remove(uuid);
                    activeRequests.remove(uuid);
                });

        // 订阅 Flux 并保存 Disposable 到 activeRequests 中
        Disposable disposable = flux.subscribeOn(Schedulers.boundedElastic())
                .doFinally(signalType -> activeRequests.remove(uuid))
                .subscribe();

        activeRequests.put(uuid, disposable);

        return flux;
    }

    // 提供一个接口用于根据 UUID 打断请求
    @GetMapping("/ai/interrupt")
    public void interruptRequest(String uuid) {
        Disposable disposable = activeRequests.get(uuid);
        AtomicBoolean isInterrupted = interruptionStatusMapV2.get(uuid);
        if (disposable != null) {
            // 标记请求为中断状态
            if (isInterrupted != null) {
                isInterrupted.set(true);
            }
            // 取消订阅，断开与 Ollama 的连接
            disposable.dispose();
            logger.info("Request with UUID {} has been interrupted.", uuid);
        }
        // 移除中断状态和 Disposable 记录
        interruptionStatusMapV2.remove(uuid);
        activeRequests.remove(uuid);
    }

    @PostMapping(value = "/ai/ollamaStreamV2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> queryModelV2(@RequestBody ChatRequest request) {
        String uuid = request.getToken();
        try {
            String requestBody = objectMapper.writeValueAsString(request);
            logger.info("Sending request body: {}", requestBody);
        } catch (Exception e) {
            logger.error("Error converting request to JSON", e);
        }

        // 初始化用户的中断状态
        AtomicBoolean isInterrupted = interruptionStatusMapV2.computeIfAbsent(uuid, k -> new AtomicBoolean(false));

        Flux<ServerSentEvent<String>> flux = webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .header("Accept-Charset", "UTF-8")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class)
                .map(response -> ServerSentEvent.<String>builder().data(response).build())
                .takeWhile(part -> !isInterrupted.get())
                .doFinally(signal -> {
                    interruptionStatusMapV2.remove(uuid);
                    activeRequests.remove(uuid);
                })
                .doOnCancel(() -> {
                    // 标记请求为中断状态
                    isInterrupted.set(true);
                    // 取消订阅，断开与 Ollama 的连接
                    Disposable disposable = activeRequests.get(uuid);
                    if (disposable != null) {
                        disposable.dispose();
                        logger.info("Request with UUID {} has been interrupted due to client cancel.", uuid);
                    }
                    // 移除中断状态和 Disposable 记录
                    interruptionStatusMapV2.remove(uuid);
                    activeRequests.remove(uuid);
                });

        // 订阅 Flux 并保存 Disposable 到 activeRequests 中
        Disposable disposable = flux.subscribeOn(Schedulers.boundedElastic())
                .doFinally(signalType -> activeRequests.remove(uuid))
                .subscribe();

        activeRequests.put(uuid, disposable);

        return flux;
    }

}
